"""
WebSocket实时数据推送服务
提供统计数据的实时更新和推送功能
"""

import json
import asyncio
from typing import Dict, List, Set, Optional
from datetime import datetime, timedelta
from fastapi import WebSocket, WebSocketDisconnect
from loguru import logger
import asyncio

from core.database import get_db
from services.petition_stats_service import PetitionStatsService
from services.multi_dimensional_stats_service import MultiDimensionalStatsService


class ConnectionManager:
    """WebSocket连接管理器"""

    def __init__(self):
        # 活跃连接管理
        self.active_connections: Dict[str, List[WebSocket]] = {}
        # 用户订阅的主题
        self.user_subscriptions: Dict[str, Set[str]] = {}
        # 数据推送间隔配置
        self.push_intervals = {
            'overview': 30,      # 概览数据30秒推送一次
            'regional': 60,     # 区域数据60秒推送一次
            'time_series': 120,  # 时间序列120秒推送一次
            'complaint_type': 90, # 投诉类型90秒推送一次
            'sentiment': 120,   # 情感分析120秒推送一次
            'compliance': 120,  # 合规性分析120秒推送一次
        }
        # 后台任务
        self.background_tasks: Dict[str, asyncio.Task] = {}

    async def connect(self, websocket: WebSocket, user_id: str):
        """建立WebSocket连接"""
        await websocket.accept()

        if user_id not in self.active_connections:
            self.active_connections[user_id] = []
            self.user_subscriptions[user_id] = set()

        self.active_connections[user_id].append(websocket)
        logger.info(f"用户 {user_id} 建立WebSocket连接，当前连接数: {len(self.active_connections[user_id])}")

    async def disconnect(self, websocket: WebSocket, user_id: str):
        """断开WebSocket连接"""
        if user_id in self.active_connections:
            if websocket in self.active_connections[user_id]:
                self.active_connections[user_id].remove(websocket)

            # 如果用户没有活跃连接了，清理相关资源
            if not self.active_connections[user_id]:
                del self.active_connections[user_id]
                if user_id in self.user_subscriptions:
                    del self.user_subscriptions[user_id]
                # 停止该用户的推送任务
                await self._stop_user_push_tasks(user_id)

        logger.info(f"用户 {user_id} 断开WebSocket连接")

    async def subscribe(self, user_id: str, topics: List[str]):
        """用户订阅主题"""
        if user_id not in self.user_subscriptions:
            self.user_subscriptions[user_id] = set()

        for topic in topics:
            self.user_subscriptions[user_id].add(topic)

        # 启动该用户的数据推送任务
        await self._start_user_push_tasks(user_id)
        logger.info(f"用户 {user_id} 订阅主题: {topics}")

    async def unsubscribe(self, user_id: str, topics: List[str]):
        """用户取消订阅主题"""
        if user_id in self.user_subscriptions:
            for topic in topics:
                self.user_subscriptions[user_id].discard(topic)

        # 如果用户没有订阅任何主题，停止推送任务
        if user_id in self.user_subscriptions and not self.user_subscriptions[user_id]:
            await self._stop_user_push_tasks(user_id)

        logger.info(f"用户 {user_id} 取消订阅主题: {topics}")

    async def send_personal_message(self, message: str, user_id: str):
        """发送个人消息"""
        if user_id in self.active_connections:
            for connection in self.active_connections[user_id]:
                try:
                    await connection.send_text(message)
                except Exception as e:
                    logger.error(f"发送消息失败: {e}")

    async def broadcast_to_topic(self, message: str, topic: str):
        """向主题订阅者广播消息"""
        for user_id, subscriptions in self.user_subscriptions.items():
            if topic in subscriptions:
                await self.send_personal_message(message, user_id)

    async def _start_user_push_tasks(self, user_id: str):
        """启动用户的数据推送任务"""
        if user_id not in self.user_subscriptions:
            return

        # 为每个订阅的主题启动推送任务
        for topic in self.user_subscriptions[user_id]:
            task_key = f"{user_id}_{topic}"
            if task_key not in self.background_tasks:
                task = asyncio.create_task(
                    self._push_data_periodically(user_id, topic)
                )
                self.background_tasks[task_key] = task

    async def _stop_user_push_tasks(self, user_id: str):
        """停止用户的数据推送任务"""
        tasks_to_stop = []
        for task_key in list(self.background_tasks.keys()):
            if task_key.startswith(f"{user_id}_"):
                task = self.background_tasks[task_key]
                if not task.done():
                    task.cancel()
                tasks_to_stop.append(task_key)

        for task_key in tasks_to_stop:
            del self.background_tasks[task_key]

    async def _push_data_periodically(self, user_id: str, topic: str):
        """定期推送数据"""
        interval = self.push_intervals.get(topic, 60)

        while user_id in self.user_subscriptions and topic in self.user_subscriptions[user_id]:
            try:
                # 准备数据
                data = await self._prepare_topic_data(topic, user_id)

                # 发送数据
                message = {
                    "type": "data_update",
                    "topic": topic,
                    "timestamp": datetime.now().isoformat(),
                    "data": data
                }

                await self.send_personal_message(json.dumps(message, ensure_ascii=False), user_id)

                # 等待下次推送
                await asyncio.sleep(interval)

            except asyncio.CancelledError:
                logger.info(f"用户 {user_id} 的 {topic} 推送任务被取消")
                break
            except Exception as e:
                logger.error(f"推送数据失败 - 用户: {user_id}, 主题: {topic}, 错误: {e}")
                await asyncio.sleep(interval)

    async def _prepare_topic_data(self, topic: str, user_id: str) -> Dict:
        """准备主题数据"""
        try:
            # 获取数据库会话
            db_generator = get_db()
            db = next(db_generator)

            # 根据主题获取不同的数据
            if topic == 'overview':
                stats_service = PetitionStatsService()
                data = await stats_service.get_statistics_overview(db)
            elif topic == 'regional':
                stats_service = PetitionStatsService()
                data = await stats_service.get_regional_statistics(db, level='district')
            elif topic == 'time_series':
                stats_service = PetitionStatsService()
                end_date = datetime.now()
                start_date = end_date - timedelta(days=30)
                data = await stats_service.get_time_series_statistics(db, start_date, end_date, 'day')
            elif topic == 'complaint_type':
                stats_service = PetitionStatsService()
                data = await stats_service.get_complaint_type_statistics(db)
            elif topic == 'sentiment':
                stats_service = PetitionStatsService()
                data = await stats_service.get_sentiment_statistics(db)
            elif topic == 'compliance':
                stats_service = PetitionStatsService()
                data = await stats_service.get_compliance_statistics(db)
            else:
                data = {"message": f"未知主题: {topic}"}

            db.close()
            return data

        except Exception as e:
            logger.error(f"准备数据失败 - 主题: {topic}, 错误: {e}")
            return {"error": str(e)}

    def get_connection_stats(self) -> Dict:
        """获取连接统计信息"""
        return {
            "total_users": len(self.active_connections),
            "total_connections": sum(len(connections) for connections in self.active_connections.values()),
            "user_subscriptions": {user_id: len(subscriptions) for user_id, subscriptions in self.user_subscriptions.items()},
            "active_topics": self._get_active_topics(),
            "background_tasks": len(self.background_tasks)
        }

    def _get_active_topics(self) -> Dict[str, int]:
        """获取活跃主题统计"""
        topic_counts = {}
        for subscriptions in self.user_subscriptions.values():
            for topic in subscriptions:
                topic_counts[topic] = topic_counts.get(topic, 0) + 1
        return topic_counts


# 全局连接管理器实例
manager = ConnectionManager()


class RealTimeDataService:
    """实时数据服务"""

    def __init__(self):
        self.manager = manager

    async def handle_websocket(self, websocket: WebSocket, user_id: str):
        """处理WebSocket连接"""
        await self.manager.connect(websocket, user_id)

        try:
            # 发送欢迎消息
            welcome_message = {
                "type": "welcome",
                "message": "连接成功",
                "timestamp": datetime.now().isoformat(),
                "available_topics": list(self.manager.push_intervals.keys())
            }
            await websocket.send_text(json.dumps(welcome_message, ensure_ascii=False))

            # 监听客户端消息
            while True:
                try:
                    # 接收客户端消息
                    data = await websocket.receive_text()
                    message = json.loads(data)

                    # 处理不同类型的消息
                    await self._handle_client_message(user_id, message)

                except json.JSONDecodeError:
                    error_message = {
                        "type": "error",
                        "message": "无效的JSON格式",
                        "timestamp": datetime.now().isoformat()
                    }
                    await websocket.send_text(json.dumps(error_message, ensure_ascii=False))

        except WebSocketDisconnect:
            await self.manager.disconnect(websocket, user_id)
        except Exception as e:
            logger.error(f"WebSocket处理错误: {e}")
            await self.manager.disconnect(websocket, user_id)

    async def _handle_client_message(self, user_id: str, message: Dict):
        """处理客户端消息"""
        message_type = message.get('type')

        if message_type == 'subscribe':
            # 订阅主题
            topics = message.get('topics', [])
            await self.manager.subscribe(user_id, topics)

            # 发送订阅确认
            response = {
                "type": "subscription_confirmed",
                "topics": topics,
                "timestamp": datetime.now().isoformat()
            }
            await self.manager.send_personal_message(json.dumps(response, ensure_ascii=False), user_id)

        elif message_type == 'unsubscribe':
            # 取消订阅
            topics = message.get('topics', [])
            await self.manager.unsubscribe(user_id, topics)

            # 发送取消订阅确认
            response = {
                "type": "unsubscription_confirmed",
                "topics": topics,
                "timestamp": datetime.now().isoformat()
            }
            await self.manager.send_personal_message(json.dumps(response, ensure_ascii=False), user_id)

        elif message_type == 'ping':
            # 心跳检测
            response = {
                "type": "pong",
                "timestamp": datetime.now().isoformat()
            }
            await self.manager.send_personal_message(json.dumps(response, ensure_ascii=False), user_id)

        elif message_type == 'get_stats':
            # 获取统计信息
            stats = self.manager.get_connection_stats()
            response = {
                "type": "connection_stats",
                "stats": stats,
                "timestamp": datetime.now().isoformat()
            }
            await self.manager.send_personal_message(json.dumps(response, ensure_ascii=False), user_id)

    async def broadcast_system_notification(self, message: str, level: str = "info"):
        """广播系统通知"""
        notification = {
            "type": "system_notification",
            "level": level,
            "message": message,
            "timestamp": datetime.now().isoformat()
        }
        await self.manager.broadcast_to_topic(json.dumps(notification, ensure_ascii=False), "system")

    async def force_refresh_data(self, topic: str):
        """强制刷新指定主题的数据"""
        refresh_message = {
            "type": "force_refresh",
            "topic": topic,
            "timestamp": datetime.now().isoformat()
        }
        await self.manager.broadcast_to_topic(json.dumps(refresh_message, ensure_ascii=False), topic)


# 全局实时数据服务实例
real_time_service = RealTimeDataService()